import json
from typing import List, Dict


def convert_to_llamafactory_format(original_data: List[Dict]) -> List[Dict]:
    """
    将原始诗歌解析数据转换为LLAMAFactory支持的格式

    Args:
        original_data: 原始数据列表，每个元素包含title/content/keywords/trans/emotion字段

    Returns:
        转换后的数据列表，符合LLAMAFactory的instruction/input/output结构
    """
    converted_data = []

    for item in original_data:
        # 构建instruction提示模板
        instruction = f"请解析这个古文《{item['title']}》诗句并生成结构化数据"

        # 结构化output数据
        structured_output = {
            "keywords": item["keywords"],
            "translation": item["trans"],
            "emotion_analysis": item["emotion"]
        }

        # 组装最终格式
        converted_item = {
            "instruction": instruction,
            "input": item["content"],
            "output": structured_output
        }

        converted_data.append(converted_item)

    return converted_data


# 示例用法
if __name__ == "__main__":
    # 原始示例数据（模拟从文件加载的数据）
    original_example = [{
        "title": "为有",
        "content": "为有云屏无限娇，凤城寒尽怕春宵。无端嫁得金龟婿，辜负香衾事早朝。",
        "keywords": {
            "云屏": "雕饰着云母图案的屏风",
            "凤城": "京城",
            "无端": "没来由",
            "金龟婿": "佩带金龟的丈夫"
        },
        "trans": "云母屏风后面的美人格外娇美...",
        "emotion": "怨恨与无奈"
    }]

    # 执行转换
    converted_data = convert_to_llamafactory_format(original_example)

    # 保存为JSON文件（LLAMAFactory推荐格式）
    # with open("llama_train_data_2.json", "w", encoding="utf-8") as f:
    #     json.dump(converted_data, f, ensure_ascii=False, indent=2)

    # 或者保存为.jsonl格式（每行一个JSON对象）
    with open("llama_train_data.jsonl", "w", encoding="utf-8") as f:
        for item in converted_data:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")

    print("转换完成！样例输出：")
    print(json.dumps(converted_data, ensure_ascii=False, indent=2))
